use anyhow::Result;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tonic::transport::{Channel, Endpoint};
use tracing::{error, info, warn};

#[derive(Debug, Clone)]
pub struct ServiceConfig {
    pub name: String,
    pub endpoints: Vec<String>,
    pub max_connections: usize,
    pub connection_timeout: Duration,
    pub request_timeout: Duration,
}

#[derive(Debug, Clone)]
pub struct ConnectionStats {
    pub total_requests: u64,
    pub failed_requests: u64,
    pub active_connections: usize,
}

pub struct ConnectionPool {
    configs: Arc<RwLock<HashMap<String, ServiceConfig>>>,
    connections: Arc<RwLock<HashMap<String, Vec<Channel>>>>,
    current_index: Arc<RwLock<HashMap<String, usize>>>,
    stats: Arc<RwLock<HashMap<String, ConnectionStats>>>,
}

impl ConnectionPool {
    pub fn new() -> Self {
        Self {
            configs: Arc::new(RwLock::new(HashMap::new())),
            connections: Arc::new(RwLock::new(HashMap::new())),
            current_index: Arc::new(RwLock::new(HashMap::new())),
            stats: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// 注册服务
    pub async fn register_service(&self, config: ServiceConfig) -> Result<()> {
        let name = config.clone().name;
        let mut configs = self.configs.write().await;
        configs.insert(name.clone(), config);

        info!("Service registered: {}", name);
        Ok(())
    }

    /// 获取连接
    pub async fn get_connection(&self, service_name: &str) -> Result<Channel, PoolError> {
        let connections = self.connections.read().await;

        if let Some(channels) = connections.get(service_name) {
            if !channels.is_empty() {
                let mut current_index = self.current_index.write().await;
                let index = current_index.entry(service_name.to_string()).or_insert(0);

                let channel = channels[*index].clone();
                *index = (*index + 1) % channels.len();

                return Ok(channel);
            }
        }

        warn!("No connections available for service: {}", service_name);
        Err(PoolError::NoConnectionsAvailable)
    }

    /// 初始化连接
    pub async fn initialize_connections(&self) -> Result<()> {
        let configs = self.configs.read().await;
        let mut connections = self.connections.write().await;
        let mut stats = self.stats.write().await;

        for (name, config) in configs.iter() {
            let mut channels = Vec::new();

            for endpoint in &config.endpoints {
                match self.create_channel(endpoint, &config).await {
                    Ok(channel) => {
                        channels.push(channel);
                        info!("Connected to {}: {}", name, endpoint);
                    }
                    Err(e) => {
                        error!("Failed to connect to {}: {} - {}", name, endpoint, e);
                    }
                }
            }

            if !channels.is_empty() {
                let len = channels.len();
                connections.insert(name.clone(), channels);
                stats.insert(
                    name.clone(),
                    ConnectionStats {
                        total_requests: 0,
                        failed_requests: 0,
                        active_connections: len,
                    },
                );
            }
        }

        Ok(())
    }

    /// 创建节点
    async fn create_channel(
        &self,
        endpoint: &str,
        config: &ServiceConfig,
    ) -> Result<Channel, tonic::transport::Error> {
        Endpoint::from_shared(endpoint.to_string())?
            .timeout(config.request_timeout)
            .connect_timeout(config.connection_timeout)
            .rate_limit(500, Duration::from_secs(1))                // 每秒最多500个请求
            .connect()
            .await
    }

    /// 获取服务配置
    pub async fn get_stats(&self, service_name: &str) -> Option<ConnectionStats> {
        let stats = self.stats.read().await;
        stats.get(service_name).cloned()
    }
}

#[derive(Debug)]
pub enum PoolError {
    NoConnectionsAvailable,
    ServiceNotFound,
    ConnectionFailed,
    RetryFailed(String),
}

impl std::fmt::Display for PoolError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            PoolError::NoConnectionsAvailable => write!(f, "No connections available"),
            PoolError::ServiceNotFound => write!(f, "Service not found"),
            PoolError::ConnectionFailed => write!(f, "Connection failed"),
            PoolError::RetryFailed(err) => write!(f, "Retry failed: {:?}", err),
        }
    }
}

impl std::error::Error for PoolError {}
